import os
import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.optimizers import Adam
from datetime import datetime
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


'''
读取CSV文件：
pd.read_csv('data_source.csv')：从CSV文件中读取数据。
将时间列转换为时间戳：
使用 datetime.strptime 将字符串形式的时间转换为时间戳。
分离特征和标签：
X 表示特征数据，y 表示标签数据。
数据标准化：
使用 StandardScaler 对特征数据进行标准化处理，使得特征具有零均值和单位方差。
划分训练集和测试集：
使用 train_test_split 函数将数据划分为训练集和测试集。
构建模型：
创建一个简单的线性回归模型，包含一个全连接层，输出单元为1，使用L2正则化。
编译模型：
使用 Adam 优化器和均方误差（MSE）作为损失函数。
打印每个批次的进度条和损失值：
自定义一个回调函数 PrintCallback，在每个epoch结束时打印训练损失。
设置训练参数：
定义训练轮次和批次大小。
学习率设置：
使用 Adam 优化器，并设置学习率为0.001。
创建回调函数：
使用 EarlyStopping 回调来防止过拟合，当验证损失在10个epoch内没有改进时停止训练。
训练模型：
使用 fit 方法训练模型，并传入训练数据、训练轮次、批次大小、进度条显示、回调函数等参数。
评估模型：
使用 evaluate 方法评估模型在测试集上的表现，并打印测试损失。
保存模型：
自动生成带有编号的文件夹名称，并保存训练好的模型。
导出模型为 SavedModel 格式：
将模型导出为 SavedModel 格式，便于后续部署或加载。
'''

# 读取CSV文件
df = pd.read_csv('data_source.csv')

# 将时间列转换为时间戳
df['time'] = df['time'].apply(lambda x: datetime.strptime(x, '%Y-%m-%d %H:%M:%S').timestamp())

# 分离特征和标签
X = df[['value', 'time']].values  # 特征数据
y = df['value'].values  # 标签数据

# 数据标准化
scaler = StandardScaler()  # 初始化StandardScaler对象
X_scaled = scaler.fit_transform(X)  # 对特征数据进行标准化

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 构建模型
model = Sequential([
    Dense(1, input_shape=(2,), activation='linear', kernel_regularizer=tf.keras.regularizers.l2(0.001))  #
])  # 创建一个具有单个输出单元的全连接层，输入为两个特征，并使用L2正则化

# 编译模型
model.compile(optimizer='adam', loss='mse')  # 使用Adam优化器和均方误差损失函数

# 打印每个批次的进度条和损失值
class PrintCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(f"\nEpoch {epoch+1}/{epochs} - Train Loss: {logs['loss']}")  # 在每个epoch结束时打印训练损失

# 设置训练参数
epochs = 50  # 训练轮次
batch_size = 2  # 批次大小

# 学习率设置
optimizer = Adam(learning_rate=0.001)  # 使用Adam优化器，学习率为0.001
model.compile(optimizer=optimizer, loss='mse')  # 再次编译模型

# 创建回调函数
callback = PrintCallback()  # 打印每个epoch结束时的训练损失
early_stopping = EarlyStopping(monitor='val_loss', patience=10)  # 如果验证损失在10个epoch内没有改进，则停止训练

# 训练模型
history = model.fit(
    X_train, y_train,  # 输入训练数据
    epochs=epochs,  # 训练轮次
    batch_size=batch_size,  # 批次大小
    verbose=1,  # 显示每个批次的进度条
    callbacks=[callback]  # 使用回调函数
)

# 评估模型
loss = model.evaluate(X_test, y_test)  # 在测试集上评估模型
print(f"Test loss: {loss}")  # 打印测试集上的损失值

# 自动生成带有编号的文件夹名称
def get_next_folder_name(base_dir):
    max_index = 0
    for entry in os.listdir(base_dir):  # 遍历基础目录下的所有条目
        if os.path.isdir(os.path.join(base_dir, entry)):  # 如果是目录
            try:
                index = int(entry.split('_')[-1])  # 获取目录名中的数字部分
                if index > max_index:  # 更新最大索引值
                    max_index = index
            except ValueError:
                continue
    return f"{base_dir}/model_{max_index + 1}"  # 返回新的目录名

# 保存模型
save_dir = get_next_folder_name('.')  # 获取下一个可用的文件夹名
os.makedirs(save_dir, exist_ok=True)  # 创建目录
model.save(os.path.join(save_dir, 'trained_model.h5'))  # 保存模型

# 可选：将模型导出为SavedModel格式
tf.saved_model.save(model, os.path.join(save_dir, 'saved_model_path'))  # 保存为 SavedModel 格式